import torch.nn as nn


class Head(nn.Module):
    def __init__(self, in_channel, num_cls: int):
        super().__init__()
        self.head_bbox = nn.Conv2d(in_channel, 4, 1)
        self.head_hm = nn.Conv2d(in_channel, num_cls, 1)
        self.head_pool = nn.MaxPool2d(3, 1, 1)

    def forward(self, x):
        out_bbox = self.head_bbox(x)
        out_hm = self.head_hm(x)
        out_pool = self.head_pool(out_hm)
        return out_bbox, out_hm, out_pool
